# %% [markdown]
# # [基线方法](https://work.datafountain.cn/forum?id=564&type=2)
# 任务描述：建立一个预测模型，通过从序列中提取特征，对microRNA和gene的序列关系进行预测数据集:microRNA的序列信息和gene信息已经提供，并且提供了一部分的关系对信息，需要使用这些关系对信息建模，预测测试集中的microRNA和gene序列是否有关系方法概述：首先从序列中提取3-mer的信息并且计数，获取microRNA和gene的3-mer的分布信息，再将二者拼接做成数据集，使用随机森林进行预测。

# %%
import pandas as pd
import itertools
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn import metrics
from sklearn.metrics import roc_auc_score
import os
import joblib

# %% [markdown]
# ## 1.数据准备
# 
# 1.1 数据说明
# 
# 1.2 数据预处理
# 
# 1.3 数据集生成

# %%
#数据说明
dataset=pd.read_csv('./datasets/train_dataset/Train.csv')
mirna_seqdf=pd.read_csv('./datasets/train_dataset/mirna_seq.csv')#(['mirna', 'seq']
gene_seqdf=pd.read_csv('./datasets/train_dataset/gene_seq.csv')#'label', 'sequence'

dataset_mirna=dataset['miRNA']
dataset_gene=dataset['gene']
dataset_label=dataset['label']
gene_index=gene_seqdf['label'].values.tolist()
gene_seq=gene_seqdf['sequence']
mirna_index=mirna_seqdf['mirna'].values.tolist()
mirna_seq=mirna_seqdf['seq']

# %% [markdown]
# ### 1.1 展示序列信息以及标签信息,以及关系对的标签信息

# %%
gene_seq.head

# %%
gene_index[0:10]

# %%
dataset_label.head

# %%
dataset.columns

# %% [markdown]
# ### 1.2 数据预处理，特征提取和生成
# A:腺嘌呤、G:鸟嘌呤、T:胸腺嘧啶、C:胞嘧啶、U:脲嘧啶。
# 
# 其中T为DNA特有，U为RNA特有。
# 
# 配对规则：A=T(双氢键)、G=C(三氢键)，RNA为单链所以U无配对。

# %%
'''
key_set与key_set_T用于保存碱基子串计数,其中key为碱基子串(字符串)、value为计数
两个循环分别用于初始化key_set与key_set_T,对其添加所有可能的碱基子串并初始化value为0
'''
# 用于碱基字串计数
key_set={}
key_set_T={}
# itertools.product :可迭代对象输入的笛卡儿积 下面这一行可以生成每个元素3元组的可迭代对象
for i in itertools.product('UCGA', repeat =3):#itertools.product('BCDEF', repeat = 2):
    #print(i)
    obj=''.join(i) #将三元组合并成一个字符串
   # print(obj)
    ky={'{}'.format(obj):0}
    key_set.update(ky)
for i in itertools.product('TCGA', repeat =3):#itertools.product('BCDEF', repeat = 2):
    #print(i)
    obj=''.join(i)
   # print(obj)
    ky={'{}'.format(obj):0}
    key_set_T.update(ky)

def clean_key_set(key_set):
    '''
    将字典所有的value设置为0
    '''
    for i,key in enumerate(key_set):
    #print(i,key,key_set[key])
      key_set[key]=0
    return key_set

# %%
def return_features(n,seq):
    '''
    n为碱基子串的长度，seq为碱基序列
    使用长度为n的滑动窗口选取碱基子串，并使用全局变量key_set记录字串数量并返回
    '''
    clean_key_set(key_set)
    key=key_set
    if '\n' in seq:
        seq=seq[0:-1]
    for i in range(n,len(seq)+1-n):#忽略前n个碱基？
        win=seq[i:i+n]#获取长度为n的碱基子串
        #print(win)
        ori=key_set['{}'.format(win)]#碱基子串计数增加
        key_set['{}'.format(win)]=ori+1
    return key_set
def return_gene_features(n,seq):
    '''
    功能可参照上一个函数
    '''
    clean_key_set(key_set_T)
    key=key_set_T
    if '\n' in seq:
        seq=seq[0:-1]
    for i in range(n,len(seq)+1-n):
        win=seq[i:i+n]
        #print(win)
        ori=key_set_T['{}'.format(win)]
        key_set_T['{}'.format(win)]=ori+1
    return key_set_T

# %% [markdown]
# ### 1.3 使用拼接方法构建数据集

# %%
def construct_dataset(dataset_mirna,dataset_gene):
    '''
    将碱基子串的计数字典作为碱基序列的特征
    将存在关系对的DNA与RNA的碱基序列特征直接连接起来作为训练用特征
    '''
    list_mirna_feature=[]
    list_gene_feature=[]
    for i in range(0,len(dataset_mirna)):# 遍历Train.csv中的RNA
        try:
            # 获取RNA
            mirna=dataset_mirna[i]#获取当前Train.csv中RNA的名字
            m_index=mirna_index.index(mirna)#根据名字获取到Train.csv中的RNA在mirna_seq.csv中的序号
            mirna_f=return_features(3,mirna_seq[m_index])#输入当前RNA的碱基序列，并获取该序列的子串计数字典
            # 获取DNA
            gene=dataset_gene[i]#获取当前Train.csv中DNA的名字
            g_index=gene_index.index(gene)#根据名字获取到Train.csv中的DNA在gene_seq.csv中的序号
            gene_f=return_gene_features(3, gene_seq[g_index])#输入当前DNA的碱基序列，并获取该序列的子串计数字典

            mirna_feature=mirna_f.copy()
            gene_feature=gene_f.copy()
            # 保存当前得到的子串计数
            list_mirna_feature.append(mirna_feature)
            list_gene_feature.append(gene_feature)
        except:
            mirna=dataset_mirna[i]
            gene=dataset_gene[i]
            print('error detected',i,mirna,gene)
    # 合并得到的DNA，RNA碱基序列特征
    lmpd=pd.DataFrame(list_mirna_feature)
    lgpd=pd.DataFrame(list_gene_feature)
    # 将对应关系的DNA，RNA序列特征直接拼接
    X=pd.concat([lmpd,lgpd],axis=1)
    return X

# %%
#标签换为数字
Y=[]
for i,label in enumerate(dataset_label):
    if label =='Functional MTI':
        Y.append(1)
    else:
        Y.append(0)

# %%
X=construct_dataset(dataset_mirna,dataset_gene)
#print(X)
#lmpd.to_csv('gene_features.csv',index=None)

# %% [markdown]
# ## 2模型训练&3.模型评估
# 模型训练切分训练集调参，使用ACC和F1score作为评估标准

# %%
X_train, X_test, y_train, y_test = train_test_split(X, Y,test_size=0.8, random_state=2)

# %%
#切分训练集进行调参
def train():
  clf = RandomForestClassifier(n_estimators=30)
  clf.fit(X_train,y_train)
  y_p=clf.predict(X_test)

  #acc = metrics.accuracy_score(y_test,y_p)
  #print('RF_ACC',acc)
  y_pb=clf.predict_proba(X_test)
  #print(y_p)
  f1score=metrics.f1_score(y_test, y_p)
  print('RF_F1',f1score)
  MCC=metrics.matthews_corrcoef(y_test, y_p)
train()

# %%
def train():
  clf = RandomForestClassifier(n_estimators=30)
  clf.fit(X_train,y_train)
  y_p=clf.predict(X_test)

  #acc = metrics.accuracy_score(y_test,y_p)
  #print('RF_ACC',acc)
  y_pb=clf.predict_proba(X_test)
  #print(y_p)
  f1score=metrics.f1_score(y_test, y_p)
  print('RF_F1',f1score)
  MCC=metrics.matthews_corrcoef(y_test, y_p)
train()

# %%
# 最终模型
clf_final = RandomForestClassifier(n_estimators=30)
clf_final.fit(X,Y)

# %%
# 存储模型与重新调用
joblib.dump(clf_final,'./model/tran_model.m')
clf_final = joblib.load('./model/tran_model.m')

# %% [markdown]
# ## 4.模型预测
# 加载测试数据
# 
# 测试数据生成器
# 
# 预测结果
# 
# 结果展示
# 
# 生成提交文件

# %%
#加载测试数据
#test_filenames = os.listdir("./datasets")
df_predict=pd.read_csv('./datasets/test_dataset.csv')

# %%
predict_mirna=df_predict['miRNA']
predict_gene=df_predict['gene']

# %%
#测试数据生成器
X_predict=construct_dataset(predict_mirna,predict_gene)

# %%
#预测结果
final_result=clf_final.predict(X_predict)

# %%
#结果展示
print(final_result)

# %%
# 生成提交文件
df_predict['results'] = final_result
df_predict.to_csv('submission.csv',index=None)


